from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
import os
import time
from crypto import AESCipher, DESCipher, KeyManager
from utils import FileHandler, Config
import base64

class HelpDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("使用帮助")
        self.setMinimumSize(600, 400)
    
        layout = QVBoxLayout(self)
        
        # 创建滚动区域
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setStyleSheet("""
            QScrollArea {
                border: none;
                background-color: transparent;
            }
        """)
        
        # 创建内容容器
        content = QWidget()
        content_layout = QVBoxLayout(content)
        
        # 添加帮助内容
        help_text = """
         <h1>1. 简介</h1>
  <p>本工具支持 AES（128/192/256位）、DES 和 3DES 算法，提供 ECB、CBC、CTR、OFB、CFB 多种加密模式，适用于文本、文件等数据的加密与解密。</p>
  <p><strong>适用场景：</strong></p>
  <ul>
    <li>保护个人隐私文件（如文档、图片）。</li>
    <li>企业敏感数据传输（如合同、数据库备份）。</li>
    <li>开发者测试加密算法性能与安全性。</li>
  </ul>

  <h1>2. 功能概览</h1>
  <h2>支持的算法与模式</h2>
  <table>
    <thead>
      <tr>
        <th>算法</th>
        <th>密钥长度</th>
        <th>推荐模式</th>
        <th>适用场景</th>
      </tr>
    </thead>
    <tbody>
      <tr>
        <td>AES-128</td>
        <td>128位（16字节）</td>
        <td>CBC、CTR</td>
        <td>通用数据加密</td>
      </tr>
      <tr>
        <td>AES-192</td>
        <td>192位（24字节）</td>
        <td>CBC、CFB</td>
        <td>高安全性需求</td>
      </tr>
      <tr>
        <td>AES-256</td>
        <td>256位（32字节）</td>
        <td>CBC、CTR</td>
        <td>军事/金融级数据保护</td>
      </tr>
      <tr>
        <td>DES</td>
        <td>56位（8字节）</td>
        <td>ECB（仅限测试）</td>
        <td>兼容旧系统</td>
      </tr>
      <tr>
        <td>3DES</td>
        <td>168位（24字节）</td>
        <td>CBC</td>
        <td>替代DES的过渡方案</td>
      </tr>
    </tbody>
  </table>

  <h2>加密模式</h2>
  <table>
    <thead>
      <tr>
        <th>加密模式</th>
        <th>特点</th>
        <th>安全建议</th>
      </tr>
    </thead>
    <tbody>
      <tr>
        <td>ECB</td>
        <td>简单快速，相同明文生成相同密文</td>
        <td>避免用于敏感数据</td>
      </tr>
      <tr>
        <td>CBC</td>
        <td>引入IV（初始化向量），安全性高</td>
        <td>推荐默认使用</td>
      </tr>
      <tr>
        <td>CTR</td>
        <td>流加密模式，支持并行处理</td>
        <td>适合实时数据流</td>
      </tr>
      <tr>
        <td>OFB</td>
        <td>错误不传播，但安全性弱于CBC</td>
        <td>慎用</td>
      </tr>
      <tr>
        <td>CFB</td>
        <td>类似CBC，但可自定义反馈大小</td>
        <td>适用于网络协议</td>
      </tr>
    </tbody>
  </table>

  <h1>3. 界面说明</h1>
  <ul>
    <li>① 算法选择：下拉菜单选择 AES-128/192/256、DES、3DES。</li>
    <li>② 模式选择：选择 ECB、CBC、CTR、OFB、CFB。</li>
    <li>③ 密钥管理：支持手动输入（16进制/Base64）或点击"生成密钥"，点击"保存密钥"可保存生成密钥。</li>
    <li>④ IV（初始化向量）管理：支持随机生成 IV 和保存 IV。文件/文本切换：输入加密文本或选择加密文件。</li>
    <li>⑤ 操作按钮：执行加密/解密。</li>
  </ul>

  <h1>4. 安全注意事项</h1>
  <h2>模式选择：</h2>
  <ul>
    <li>禁止使用 ECB 模式加密敏感数据（易被模式分析攻击）。</li>
    <li>推荐组合：<strong>AES-256 + CBC/CTR</strong> 模式。</li>
  </ul>

  <h2>IV（初始化向量）管理：</h2>
  <ul>
    <li>CBC、CFB、OFB、CTR 等模式必须使用 IV，ECB 模式无需 IV。</li>
    <li>解密时需要提供相同的 IV。（本工具将 IV 和密钥一同存储）</li>
  </ul>

  <h2>文件加密警告：</h2>
  <ul>
    <li>加密后立即删除原始文件，防止恢复。</li>
    <li>大文件（&gt;1GB）加密时确保磁盘空间充足。</li>
  </ul>
        """
        
        help_label = QLabel(help_text)
        help_label.setWordWrap(True)
        help_label.setTextFormat(Qt.RichText)
        content_layout.addWidget(help_label)
        
        scroll.setWidget(content)
        layout.addWidget(scroll)
        
        # 添加关闭按钮
        close_btn = QPushButton("关闭")
        close_btn.clicked.connect(self.accept)
        layout.addWidget(close_btn, alignment=Qt.AlignCenter)

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("加密解密工具")
        self.setMinimumSize(800, 600)
        
        # 初始化组件
        self.config = Config()
        self.key_manager = KeyManager(self.config.get('key_file'))
        self.file_handler = FileHandler()
        
        self.init_ui()
        self.setup_menu()
        
    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 加密设置组
        settings_group = QGroupBox("加密设置")
        settings_layout = QGridLayout()
        
        # 算法选择
        self.algo_combo = QComboBox()
        self.algo_combo.addItems(["AES-128", "AES-192", "AES-256", "DES", "3DES"])
        settings_layout.addWidget(QLabel("算法:"), 0, 0)
        settings_layout.addWidget(self.algo_combo, 0, 1)
        
        # 模式选择
        self.mode_combo = QComboBox()
        self.mode_combo.addItems(["CBC", "CTR", "OFB", "CFB","ECB"])
        settings_layout.addWidget(QLabel("模式:"), 0, 2)
        settings_layout.addWidget(self.mode_combo, 0, 3)
        
        settings_group.setLayout(settings_layout)
        layout.addWidget(settings_group)
        
        # 密钥管理组
        key_group = QGroupBox("密钥管理")
        key_layout = QHBoxLayout()
        
        self.key_input = QLineEdit()
        self.key_input.setEchoMode(QLineEdit.Password)
        self.key_input.setPlaceholderText("输入密钥或选择保存的密钥")
        
        self.key_select = QComboBox()
        self.refresh_key_list()
        
        generate_key_btn = QPushButton("生成密钥")
        save_key_btn = QPushButton("保存密钥")
        
        key_layout.addWidget(self.key_input)
        key_layout.addWidget(self.key_select)
        key_layout.addWidget(generate_key_btn)
        key_layout.addWidget(save_key_btn)
        
        key_group.setLayout(key_layout)
        layout.addWidget(key_group)
        
        # IV管理组
        iv_group = QGroupBox("IV管理")
        iv_layout = QHBoxLayout()
        
        self.iv_input = QLineEdit()
        self.iv_input.setPlaceholderText("输入IV或选择保存的IV")
        
        self.iv_select = QComboBox()
        self.refresh_iv_list()
        
        generate_iv_btn = QPushButton("生成IV")
        save_iv_btn = QPushButton("保存IV")
        
        iv_layout.addWidget(self.iv_input)
        iv_layout.addWidget(self.iv_select)
        iv_layout.addWidget(generate_iv_btn)
        iv_layout.addWidget(save_iv_btn)
        
        iv_group.setLayout(iv_layout)
        layout.addWidget(iv_group)
        
        # 文本输入/输出组
        text_group = QGroupBox("文本处理")
        text_layout = QVBoxLayout()
        
        self.input_text = QTextEdit()
        self.output_text = QTextEdit()
        self.input_text.setPlaceholderText("输入要处理的文本...")
        self.output_text.setPlaceholderText("输出结果...")
        self.output_text.setReadOnly(True)
        
        text_layout.addWidget(self.input_text)
        text_layout.addWidget(self.output_text)
        
        text_group.setLayout(text_layout)
        layout.addWidget(text_group)
        
        # 操作按钮
        button_layout = QHBoxLayout()
        self.encrypt_btn = QPushButton("加密")
        self.decrypt_btn = QPushButton("解密")
        self.file_btn = QPushButton("选择文件")
        self.help_btn = QPushButton("帮助")
                
        button_layout.addWidget(self.encrypt_btn)
        button_layout.addWidget(self.decrypt_btn)
        button_layout.addWidget(self.file_btn)
        button_layout.addWidget(self.help_btn)
        
        layout.addLayout(button_layout)
        
        # 连接信号
        self.encrypt_btn.clicked.connect(self.encrypt_data)
        self.decrypt_btn.clicked.connect(self.decrypt_data)
        self.file_btn.clicked.connect(self.select_file)
        self.help_btn.clicked.connect(self.show_help)
        generate_key_btn.clicked.connect(self.generate_key)
        save_key_btn.clicked.connect(self.save_current_key)
        self.key_select.currentTextChanged.connect(self.load_selected_key)
        generate_iv_btn.clicked.connect(self.generate_iv)
        save_iv_btn.clicked.connect(self.save_current_iv)
        self.iv_select.currentTextChanged.connect(self.load_selected_iv)
        
        # 添加模式切换时的IV处理
        self.mode_combo.currentTextChanged.connect(self.on_mode_changed)
        self.algo_combo.currentTextChanged.connect(self.on_algo_changed)
        
    def setup_menu(self):
        menubar = self.menuBar()
        
        # 文件菜单
        file_menu = menubar.addMenu('文件')
        open_action = QAction('打开...', self)
        save_action = QAction('保存...', self)
        exit_action = QAction('退出', self)
        
        file_menu.addAction(open_action)
        file_menu.addAction(save_action)
        file_menu.addSeparator()
        file_menu.addAction(exit_action)
        
        # 连接菜单动作
        open_action.triggered.connect(self.select_file)
        save_action.triggered.connect(self.save_output)
        exit_action.triggered.connect(self.close)
    
    def refresh_key_list(self):
        """刷新密钥选择下拉框。"""
        self.key_select.clear()
        self.key_select.addItem("-- 选择已保存密钥 --")
        self.key_select.addItems(self.key_manager.list_keys())
    
    def refresh_iv_list(self):
        """刷新IV选择下拉框。"""
        self.iv_select.clear()
        self.iv_select.addItem("-- 选择已保存IV --")
        self.iv_select.addItems(self.key_manager.list_ivs())
    
    def get_cipher(self):
        """根据当前设置获取适当的加密器。"""
        key = self.key_input.text().encode()
        algo = self.algo_combo.currentText()
        mode = self.mode_combo.currentText()
        
        # 获取IV/nonce（如果需要）
        iv = None
        nonce = None
        if mode != "ECB":
            iv_text = self.iv_input.text()
            if iv_text:
                try:
                    if mode == "CTR":
                        nonce = base64.b64decode(iv_text)
                        # 验证nonce长度
                        expected_size = 12 if algo.startswith("AES") else 7
                        if len(nonce) != expected_size:
                            QMessageBox.warning(self, "警告", 
                                f"{algo} CTR模式要求nonce长度为{expected_size}字节，当前长度为{len(nonce)}字节")
                            return None
                    else:
                        iv = base64.b64decode(iv_text)
                        # 验证IV长度
                        expected_size = 16 if algo.startswith("AES") else 8
                        if len(iv) != expected_size:
                            QMessageBox.warning(self, "警告", 
                                f"{algo}加密要求IV长度为{expected_size}字节，当前IV长度为{len(iv)}字节")
                            return None
                except Exception:
                    QMessageBox.warning(self, "警告", "IV/nonce格式不正确，请使用Base64编码")
                    return None
        
        if algo.startswith("AES"):
            key_size = int(algo.split("-")[1]) // 8
            mode_map = {
                "ECB": AESCipher.MODE_ECB,
                "CBC": AESCipher.MODE_CBC,
                "CTR": AESCipher.MODE_CTR,
                "OFB": AESCipher.MODE_OFB,
                "CFB": AESCipher.MODE_CFB
            }
            try:
                cipher = AESCipher(key[:key_size], mode=mode_map[mode], iv=iv, nonce=nonce)
            except ValueError as e:
                QMessageBox.warning(self, "警告", str(e))
                return None
        else:
            mode_map = {
                "ECB": DESCipher.MODE_ECB,
                "CBC": DESCipher.MODE_CBC,
                "CTR": DESCipher.MODE_CTR,
                "OFB": DESCipher.MODE_OFB,
                "CFB": DESCipher.MODE_CFB
            }
            try:
                cipher = DESCipher(key, triple_des=algo == "3DES", mode=mode_map[mode], iv=iv, nonce=nonce)
            except ValueError as e:
                QMessageBox.warning(self, "警告", str(e))
                return None
        return cipher
    
    def encrypt_data(self):
        """加密输入数据。"""
        try:
            cipher = self.get_cipher()
            if not cipher:
                return
                
            input_data = self.input_text.toPlainText()
            if not input_data:
                QMessageBox.warning(self, "警告", "请输入要加密的内容")
                return
            
            # 记录开始时间
            start_time = time.perf_counter()
            
            encrypted, iv = cipher.encrypt(input_data)
            
            # 计算耗时
            elapsed_time = time.perf_counter() - start_time
            
            # 格式化时间显示
            if elapsed_time < 1:
                time_str = f"{elapsed_time*1000:.2f}毫秒"
            else:
                time_str = f"{elapsed_time:.2f}秒"
            
            # 在输出文本中显示加密结果和耗时
            self.output_text.setText(f"加密结果：\n{encrypted}\n\n耗时：{time_str}")
            
            # 显示给用户
            if iv:
                self.iv_input.setText(iv)
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"加密失败: {str(e)}")
    
    def decrypt_data(self):
        """解密输入数据。"""
        try:
            cipher = self.get_cipher()
            if not cipher:
                return
                
            input_data = self.input_text.toPlainText()
            if not input_data:
                QMessageBox.warning(self, "警告", "请输入要解密的内容")
                return
                
            # 获取IV/nonce
            iv = None
            if self.mode_combo.currentText() != "ECB":
                iv = self.iv_input.text()
                if not iv:
                    QMessageBox.warning(self, "警告", "非ECB模式需要提供IV/nonce")
                    return
            
     
            start_time = time.perf_counter()
            
            decrypted = cipher.decrypt(input_data, iv)

            elapsed_time = time.perf_counter() - start_time
            
            # 格式化时间显示
            if elapsed_time < 1:
                time_str = f"{elapsed_time*1000:.2f}毫秒"
            else:
                time_str = f"{elapsed_time:.2f}秒"
            
            # 在输出文本中显示解密结果和耗时
            self.output_text.setText(f"解密结果：\n{decrypted}\n\n耗时：{time_str}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"解密失败: {str(e)}")
    
    def select_file(self):
        """打开文件对话框并处理选中的文件。"""
        file_path, _ = QFileDialog.getOpenFileName(self, "选择文件")
        if file_path:
            try:
                content, file_type = self.file_handler.read_file(file_path)
                if file_type == 'text':
                    self.input_text.setText(content.decode('utf-8'))
                else:
                    self.input_text.setText(base64.b64encode(content).decode('utf-8'))
                self.config.add_recent_file(file_path)
            except Exception as e:
                QMessageBox.critical(self, "错误", f"读取文件失败: {str(e)}")
    
    def save_output(self):
        """将输出文本保存到文件。"""
        if not self.output_text.toPlainText():
            QMessageBox.warning(self, "警告", "没有可保存的内容")
            return
            
        file_path, _ = QFileDialog.getSaveFileName(self, "保存文件")
        if file_path:
            try:
                with open(file_path, 'w') as f:
                    f.write(self.output_text.toPlainText())
                QMessageBox.information(self, "成功", "文件保存成功")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存文件失败: {str(e)}")
    
    def generate_key(self):
        """生成并设置新密钥。"""
        algo = self.algo_combo.currentText()
        if algo.startswith("AES"):
            key_size = int(algo.split("-")[1]) // 8
        else:
            key_size = 24 if algo == "3DES" else 8
            
        key = self.key_manager.generate_key(key_size=key_size)
        self.key_input.setText(base64.b64encode(key).decode())
    
    def save_current_key(self):
        """保存当前密钥。"""
        key = self.key_input.text()
        if not key:
            QMessageBox.warning(self, "警告", "请先输入或生成密钥")
            return
            
        name, ok = QInputDialog.getText(self, "保存密钥", "请输入密钥名称:")
        if ok and name:
            try:
                self.key_manager.save_key(name, key.encode())
                self.refresh_key_list()
                QMessageBox.information(self, "成功", "密钥保存成功")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存密钥失败: {str(e)}")
    
    def load_selected_key(self, key_name):
        """从存储中加载选中的密钥。"""
        if key_name and key_name != "-- 选择已保存密钥 --":
            try:
                key = self.key_manager.load_key(key_name)
                self.key_input.setText(base64.b64encode(key).decode())
            except Exception as e:
                QMessageBox.critical(self, "错误", f"加载密钥失败: {str(e)}")
    
    def generate_iv(self):
        """生成并设置新IV/nonce。"""
        algo = self.algo_combo.currentText()
        mode = self.mode_combo.currentText()
        
        # 如果是ECB模式，不需要IV
        if mode == "ECB":
            QMessageBox.information(self, "提示", "ECB模式不需要IV")
            return
            
        # 根据算法和模式确定长度
        if mode == "CTR":
            size = 12 if algo.startswith("AES") else 7  # CTR模式使用nonce
        else:
            size = 16 if algo.startswith("AES") else 8  # 其他模式使用IV
            
        iv = self.key_manager.generate_iv(iv_size=size)
        self.iv_input.setText(base64.b64encode(iv).decode())
    
    def save_current_iv(self):
        """保存当前IV。"""
        iv = self.iv_input.text()
        if not iv:
            QMessageBox.warning(self, "警告", "请先输入或生成IV")
            return
            
        name, ok = QInputDialog.getText(self, "保存IV", "请输入IV名称:")
        if ok and name:
            try:
                self.key_manager.save_iv(name, iv.encode())
                self.refresh_iv_list()
                QMessageBox.information(self, "成功", "IV保存成功")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存IV失败: {str(e)}")
    
    def load_selected_iv(self, iv_name):
        """从存储中加载选中的IV。"""
        if iv_name and iv_name != "-- 选择已保存IV --":
            try:
                iv = self.key_manager.load_iv(iv_name)
                self.iv_input.setText(base64.b64encode(iv).decode())
            except Exception as e:
                QMessageBox.critical(self, "错误", f"加载IV失败: {str(e)}")
    
    def show_help(self):
        """显示帮助对话框"""
        dialog = HelpDialog(self)
        dialog.exec_()
    
    def on_mode_changed(self, mode):
        """当加密模式改变时的处理。"""
        if mode == "ECB":
            self.iv_input.clear()
            self.iv_input.setEnabled(False)
            self.iv_select.setEnabled(False)
        else:
            self.iv_input.setEnabled(True)
            self.iv_select.setEnabled(True)
            # 如果当前有IV/nonce，验证其长度
            if self.iv_input.text():
                try:
                    iv = base64.b64decode(self.iv_input.text())
                    if mode == "CTR":
                        expected_size = 12 if self.algo_combo.currentText().startswith("AES") else 7
                    else:
                        expected_size = 16 if self.algo_combo.currentText().startswith("AES") else 8
                    if len(iv) != expected_size:
                        self.iv_input.clear()
                        QMessageBox.warning(self, "警告", 
                            f"当前IV/nonce长度与{mode}模式不匹配，已清空输入框")
                except Exception:
                    self.iv_input.clear()
    
    def on_algo_changed(self, algo):
        """当加密算法改变时的处理。"""
        # 如果当前有IV/nonce，验证其长度
        if self.iv_input.text():
            try:
                iv = base64.b64decode(self.iv_input.text())
                mode = self.mode_combo.currentText()
                if mode == "CTR":
                    expected_size = 12 if algo.startswith("AES") else 7
                else:
                    expected_size = 16 if algo.startswith("AES") else 8
                if len(iv) != expected_size:
                    self.iv_input.clear()
                    QMessageBox.warning(self, "警告", 
                        f"当前IV/nonce长度与{algo}不匹配，已清空输入框")
            except Exception:
                self.iv_input.clear()